import os
import re

# 输入输出目录
INPUT_DIR = "samples/subtrain/asm/train"
OUTPUT_DIR = "samples/subtrain/clean/train"

os.makedirs(OUTPUT_DIR, exist_ok=True)  # 创建输出文件夹

def clean_line(line):
    # 去掉地址和机器码
    line = re.sub(r'^[^a-z]*', '', line)
    # 合并多个空格为一个
    line = re.sub(r'\s+', ' ', line).strip()
    # 删除行中 ; 后的注释
    line = re.sub(r';.*$', '', line).strip()
    #合并多个下划线为一个
    line = re.sub(r'_+', '_', line)
    return line

for file_name in os.listdir(INPUT_DIR):
    if not file_name.lower().endswith(".asm"):
        continue

    input_file = os.path.join(INPUT_DIR, file_name)
    output_file = os.path.join(OUTPUT_DIR, os.path.splitext(file_name)[0] + ".asm")

    cleaned_instructions = []
    with open(input_file, 'r', encoding='utf-8', errors='ignore') as f:
        for line in f:
            line = line.split(';')[0].strip()
            if not line:
                continue
            if line.startswith(".text:") or line.startswith("CODE:") or line.startswith("code:") or line.startswith(".edata:") or line.startswith(".code:"):
                line = line.split(":", 1)[1].strip()
            else:
                continue
            line = clean_line(line)
            line = line.lower()
            if line:
                cleaned_instructions.append(line)

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write('\n'.join(cleaned_instructions))

    print(f"处理完成：{file_name} → {output_file}")

print("所有文件处理完成！")
